package sm4fast

import (
	"fmt"

	"github.com/tjfoc/gmsm/sm4"
)

// SM4ECBFast 高性能 SM4 ECB 实现
type SM4ECBFast struct {
	encryptBlocks func(dst, src []byte)
	decryptBlocks func(dst, src []byte)
	blockSize     int
}

// NewSM4ECBFast 创建高性能 SM4 ECB 实例
func NewSM4ECBFast(key []byte) (*SM4ECBFast, error) {
	if len(key) != sm4.BlockSize {
		return nil, fmt.Errorf("invalid key size")
	}

	block, err := sm4.NewCipher(key)
	if err != nil {
		return nil, err
	}

	return &SM4ECBFast{
		encryptBlocks: block.Encrypt,
		decryptBlocks: block.Decrypt,
		blockSize:     block.BlockSize(),
	}, nil
}

// Encrypt 加密
func (s *SM4ECBFast) Encrypt(plaintext []byte) []byte {
	plaintext = s.pkcs5Padding(plaintext)
	ciphertext := make([]byte, len(plaintext))

	for i := 0; i < len(plaintext); i += s.blockSize {
		s.encryptBlocks(ciphertext[i:], plaintext[i:])
	}

	return ciphertext
}

// Decrypt 解密
func (s *SM4ECBFast) Decrypt(ciphertext []byte) ([]byte, error) {
	if len(ciphertext)%s.blockSize != 0 {
		return nil, fmt.Errorf("invalid ciphertext length")
	}

	plaintext := make([]byte, len(ciphertext))

	for i := 0; i < len(ciphertext); i += s.blockSize {
		s.decryptBlocks(plaintext[i:], ciphertext[i:])
	}

	return s.pkcs5Unpadding(plaintext)
}

// pkcs5Padding 填充
func (s *SM4ECBFast) pkcs5Padding(src []byte) []byte {
	padding := s.blockSize - len(src)%s.blockSize
	padtext := make([]byte, padding)
	for i := range padtext {
		padtext[i] = byte(padding)
	}
	return append(src, padtext...)
}

// pkcs5Unpadding 去除填充
func (s *SM4ECBFast) pkcs5Unpadding(src []byte) ([]byte, error) {
	length := len(src)
	unpadding := int(src[length-1])

	if unpadding > length {
		return nil, fmt.Errorf("invalid padding")
	}

	return src[:length-unpadding], nil
}
